{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2404c024",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nuscenes.nuscenes import NuScenes\n",
    "from ultralytics import YOLO, SAM\n",
    "import torch\n",
    "import numpy as np\n",
    "import cv2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5fda599",
   "metadata": {},
   "outputs": [],
   "source": [
    "det_model = YOLO('./ckpts/yolo11l-seg.pt')\n",
    "sam_model = SAM('./ckpts/sam2.1_l.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "782fa7fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "nusc = NuScenes(version='v1.0-trainval-select', dataroot='../data/nuscenes', verbose=True)\n",
    "sample_token = nusc.scene[3]['first_sample_token']\n",
    "sample = nusc.get('sample', sample_token)\n",
    "# cams = ['CAM_FRONT', 'CAM_FRONT_RIGHT', 'CAM_BACK_RIGHT', 'CAM_BACK', 'CAM_BACK_LEFT', 'CAM_FRONT_LEFT']\n",
    "cams = ['CAM_FRONT']\n",
    "cam_token_list =  [sample['data'][cam] for cam in cams]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5ac588b",
   "metadata": {},
   "outputs": [],
   "source": [
    "fps = 10\n",
    "w, h = 1600, 900\n",
    "video_writer_list =  []\n",
    "for cam in cams:\n",
    "    video_writer_list.append(cv2.VideoWriter(f\"isegment_output_{cam}.mp4\", cv2.VideoWriter_fourcc(*\"mp4v\"), fps, (w, h)))\n",
    "# cv2.VideoWriter(\"isegment_output.mp4\", cv2.VideoWriter_fourcc(*\"mp4v\"), fps, (w, h))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9d24465",
   "metadata": {},
   "outputs": [],
   "source": [
    "count = 0\n",
    "while cam_token_list[0] != '':\n",
    "    data_path_list = []\n",
    "    for i in range(len(cams)):\n",
    "        cam = nusc.get('sample_data', cam_token_list[i])\n",
    "        data_path_list.append('../data/nuscenes/' + cam['filename'])\n",
    "        cam_token_list[i] = cam['next']\n",
    "\n",
    "    det_results = det_model.track(data_path_list, persist=True)\n",
    "    for i in range(len(cams)):\n",
    "        # class_ids = det_results[i].boxes.cls.int().tolist()\n",
    "        # if class_ids:\n",
    "        #     boxes = det_results[i].boxes.xyxy  # Boxes object for bbox outputs\n",
    "        #     sam_results = sam_model(det_results[i].orig_img, bboxes=boxes, labels=class_ids, save=False)\n",
    "        #     det_results[i].masks = sam_results[0].masks\n",
    "        video_writer_list[i].write(det_results[i].plot(line_width=2, font_size=2, conf=False))\n",
    "\n",
    "    \n",
    "    # class_ids = det_results[0].boxes.cls.int().tolist()  # Extract class IDs from detection results\n",
    "    # if class_ids:\n",
    "    #     boxes = det_results[0].boxes.xyxy  # Boxes object for bbox outputs\n",
    "    #     sam_results = sam_model(det_results[0].orig_img, bboxes=boxes, labels=class_ids, save=False)\n",
    "    # det_results[0].masks = sam_results[0].masks\n",
    "    # video_writer.write(det_results[0].plot(line_width=2, font_size=2, conf=False))\n",
    "\n",
    "    # count += 1\n",
    "    # if count >= 20:\n",
    "    #     break\n",
    "\n",
    "for video_writer in video_writer_list:\n",
    "    video_writer.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61b76835",
   "metadata": {},
   "outputs": [],
   "source": [
    "video_writer.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b72a2cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "for video_writer in video_writer_list:\n",
    "    video_writer.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1242d134",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python310",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
